from model import *
from utils import *
import random

mnist_test = torchvision.datasets.MNIST(root="./data",
                                        train=False,
                                        transform=transforms.ToTensor(),
                                        download=True)

model = VAE()
model.load_state_dict(torch.load('vae.params'))
model.eval()

show_images(model, mnist_test, random.randint(0, len(mnist_test) - 5))
